
from __future__ import annotations

from enum import Enum
from typing import List, Optional, Union

from pydantic import BaseModel, Field, conint


class ValidationError(BaseModel):
    loc: List[Union[str, int]] = Field(..., title='Location')
    msg: str = Field(..., title='Message')
    type: str = Field(..., title='Error Type')


class HTTPValidationError(BaseModel):
    detail: Optional[List[ValidationError]] = Field(None, title='Detail')


class LigandFormat(Enum):
    mol2 = 'mol2'
    sdf = 'sdf'
    txt = 'txt'


class MolecularDockingRequest(BaseModel):
    protein: Optional[str] = Field(None, title='Protein')
    ligand: Optional[str] = Field(None, title='Ligand')
    ligand_file_type: Optional[LigandFormat] = None
    num_poses: Optional[conint(le=100, gt=0)] = Field(
        10, title='Number of Poses to Generate'
    )
    time_divisions: Optional[conint(le=20, gt=3)] = Field(
        20, title='Number of Time Divisions'
    )
    steps: Optional[conint(le=18, gt=1)] = Field(18, title='Number of Diffusion Steps')
    save_trajectory: Optional[bool] = Field(False, title='Save trajectory?')
    is_staged: Optional[bool] = Field(False, title='Is staged?')
